Amazon SageMakerでマネージドなMLflowが使用できるようになりました

nokomoro3
2024.06.20
こんちには。
データアナリティクス事業本部 機械学習チームの中村(nokomoro3)です。
今朝目が覚めたら、Amazon SageMakerでマネージドなMLflowが使用可能になっていました。
アナウンス
Amazon SageMaker now offers a fully managed MLflow Capability - AWS

公式ブログ
Announcing the general availability of fully managed MLflow on Amazon SageMaker | AWS News Blog

ドキュメント
Manage machine learning experiments using Amazon SageMaker with MLflow - Amazon SageMaker

料金
Machine Learning Service – Amazon SageMaker Pricing – AWS

SageMaker Studioを立ち上げると、ApplicationsのところにMLflowが追加されています。
今回はこちらの機能を試してみたいと思います。
 MLflowを立ち上げてみる 確認先ほどのApplicationsからMLflowを選択すると以下のような画面になりますので、「Create」を押下します。
必要項目を入力します。
説明文を読んでみましょう。
By default, MLflow Tracking Servers will be size Small, use MLflow 2.13.2, and use the execution role for the SageMaker Studio Domain.

Please ensure you have configured the appropriate IAM role for this MLflow Tracking Server. By default we will use the Studio Domain execution role as the tracking server role. For more details, please see the SageMaker documentation.
デフォルトではMLflowのTracking ServerのサイズはSmall、MLflowのバージョンは 2.13.2 、ロールはドメインの実行ロールが使われるようです。

（今回はSageMakerドメインの高速セットアップで自動的に作成された実行ロールとなります）
これらの設定値は「Update」を押下すると編集することができます。値については以下にも説明がありました。
Manage machine learning experiments using Amazon SageMaker with MLflow - Amazon SageMaker
Tracking Serverのサイズは、要求されるTPS（Transactions Per Second）によってSmall、Medium、Largeから選択します。
MLflowのバージョンは今は 2.13.2 のみサポートしているようです。
実行ロールが適切な権限をもっているかどうかは、以下で確認できそうです。
Set up IAM permissions for MLflow - Amazon SageMaker
今回はそのまま作成を押下します。するとExperimentsの画面に遷移します。
MLflow自体の起動状況は、MLflowの画面で以下のように確認できます。
しばらくまつと「Creating」から「On」に遷移します。
「Open MLflow」を押下すると、MLflowの画面を確認することができます。
ただしこのままではアクセスできませんでした。現在のところは、高速セットアップで作成したドメインの実行ロールには、必要な権限が含まれていなかったようです。
 IAMポリシーの修正高速セットアップで作成したドメインの実行ロールには以下のポリシーが割り当たっています。
マネージドポリシー
AmazonSageMakerCanvasAIServicesAccess
AmazonSageMakerCanvasFullAccess
AmazonSageMakerFullAccess

カスタムポリシー
AmazonSageMaker-ExecutionPolicy-{タイムスタンプ}

この内容を以下と比較して、不足があることを確認できました。
Set up IAM permissions for MLflow - Amazon SageMaker
不足分は sagemaker-mlflow:* の部分の権限だったので、以下のように修正しました。
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "s3:GetObject",
                "s3:PutObject",
                "s3:DeleteObject",
                "s3:ListBucket"
            ],
            "Resource": [
                "arn:aws:s3:::*"
            ]
        },
        {            
            "Effect": "Allow",            
            "Action": [
                "sagemaker-mlflow:*"
            ],            
            "Resource": "*"        
        }
    ]
}
 再確認カスタムポリシー修正、以下のようにMLflowの画面にアクセスできることが確認できました。
 実験管理を試す サンプルノートブックをJupyterLab Spaceで立ち上げる今度はExperimentsの画面から、「Open sample notebook」を押下して実験を実行してみましょう。
サンプルのプレビュー画面となりますので、さらに「Open sample notebook」を押下します。
ノートブックを実行するためのJupyterLab Spaceを選択するダイアログが表示されます。今回は新しいSpaceを作成して作成ボタンを押下します。
以下のようにStatusが「Running」となるまで待ちます。
このままJupyterLabを開いてもサンプルのノートブックは現れませんでしたので、再度Experimentsの画面から「Open sample notebook」のダイアログを表示して、今度は「Use Existing Notebook」で先ほど作成したJupyterLab Spaceを選択して起動します。
すると、JupyterLab Spaceでサンプルのノートブックが開かれた状態に遷移できました。
 MLflowのTracking Serverの設定ノートブックの <TRACKING_SERVER_ARN> の部分は編集する必要があります。事前に作成したMLflowのリソースからARNを複製します。
そしてノートブック内の以下の値を差し替えます。
# Provide the ARN of the Tracking Server that you want to track your training job with
tracking_server_arn = '<TRACKING_SERVER_ARN>'
 サンプルノートブックの修正その他、いくつかそのままでは動かない部分があるため対応を実施します。
 SageMaker SDKを最新化まずはSageMaker SDKを最新化します。（切り分けの過程でこちらが必須かどうかまでは不明ですが年のため）
起動されるJupyterLab Space次第ではありますが、今回は sagemaker-distribution:1.8.0-cpu でJupyterLab Spaceが作成されたため、SageMaker SDKのバージョンは 2.219.0 となっていました。こちらを以下のコマンドで最新の 2.224.0 にアップグレードします。
Release v2.224.0 · aws/sagemaker-python-sdk
!pip install --upgrade sagemaker
バージョンを確認するには、以下を実行ください。最新になっていないようであればノートブックを再起動しましょう。
import sagemaker
print(sagemaker.__version__)

# sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml
# sagemaker.config INFO - Not applying SDK defaults from location: /home/sagemaker-user/.config/sagemaker/config.yaml
# 2.224.0
 requirements.txtの修正追加でインストールするPythonパッケージを管理するrequirements.txtが、サンプルノートブックではたくさん記載されていますが、mlflowとsagemaker-mlflowだけにしておきます。

（サンプルノートブックの通り全部をそのまま使うと、指定のバージョンが見つからないケースがありました）
%%writefile training_code/requirements.txt
mlflow==2.13.2
sagemaker-mlflow==0.1.0
 warm poolの無効化最後に以下で指定されている keep_alive_period_in_seconds=3600 を削除またはコメントアウトします。
sklearn = SKLearn(
    entry_point='train.py',
    source_dir='training_code',
    framework_version='1.2-1',
    instance_type='ml.c4.xlarge',
    role=role,
    sagemaker_session=sagemaker_session,
    hyperparameters={'max_leaf_nodes': 30},
    # keep_alive_period_in_seconds=3600,
    environment={
        'MLFLOW_TRACKING_ARN': tracking_server_arn
    }
)
こちらが有効になっていると、通常のTrainingジョブとは別のサービスクォータ「for training warm pool usage」の引き上げが必要になるため、設定を削除して「for training job usage」側の制限が適用されるようにしました。サービスクォータの設定値によってはそのままでも問題ないはずです。
 ノートブックの実行これらの修正が終わった後、ノートブックのすべてのセルを一括実行しれば完了です。
 結果の確認Trainingジョブが完了すると、MLflowの画面が更新され、以下のように学習結果が表示されることを確認できました。
いろいろクリックして学習結果やメトリクスを確認できますので、触ってみてください。
 モデル登録MLflowの画面からModelを登録することが可能です。以下のように「Run Name」のところをクリックします。
右上の「Register model」をクリックします。
更新する場合はモデルを選択すればよいはずですが、今回は新規に「Create New Model」を選択します。
Model Nameを入力して、「Register」を押下します。
Detailsのところに、Registered Modelが追加されていることが分かると思います。
SageMaker StudioのModelsからも、登録したモデルを確認することができました。
 リソースの停止料金が発生しないよう、使用したリソースを停止していきます。
まずは「Running instances」から、JupyterLab Spaceを停止します。
またMLflowのページから「Delete」もしくは「Stop」を実行します。
Deleteの場合は、実験管理のメタデータ（メトリクスやパラメータ等）が失われるのでご注意ください。（学習したモデルのデータはS3に格納されているので失われないようです）
Are you sure you want to delete the tracking server cm-nakamura-mlflow?

Deleting a tracking server will result in losing all experiment metadata (metrics, parameters, runs) you logged to it. Your artifacts (model weights, images, datasets) will still remain in the S3 bucket configured for this tracking server.
 まとめいかがでしたでしょうか。本記事がSageMaker Studioで使用できるようになったマネージドなMLflowが気になる方への参考になれば幸いです。